package cn.yx.common.mdc.wrapper;

import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.Callable;

/**
 * <p>MDC包装Runnable和Callable，线程中可支持MDC上下文信息</p>
 *
 * @author Wgssmart
 */
public class MdcWrapper {

    private MdcWrapper() {
    }

    /**
     * 使用MDC包装Callable
     *
     * @param task
     * @param contextMap
     * @param <T>
     * @return
     */
    public static <T> Callable<T> wrap(final Callable<T> task, final Map<String, String> contextMap) {
        return () -> {
            if (!contextMap.isEmpty()) {
                MDC.setContextMap(contextMap);
            }
            try {
                return task.call();
            } finally {
                if (!contextMap.isEmpty()) {
                    MDC.clear();
                }
            }
        };
    }

    /**
     * 使用MDC包装Runnable
     *
     * @param runnable
     * @param contextMap
     * @return
     */
    public static Runnable wrap(final Runnable runnable, final Map<String, String> contextMap) {
        return () -> {
            if (!contextMap.isEmpty()) {
                MDC.setContextMap(contextMap);
            }
            try {
                runnable.run();
            } finally {
                if (!contextMap.isEmpty()) {
                    MDC.clear();
                }
            }
        };
    }

}
